from autogen_agentchat.base import ChatAgent, TerminationCondition
from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, MessageFactory
from autogen_agentchat.teams import BaseGroupChat
from autogen_agentchat.teams._group_chat._events import GroupChatTermination
from pydantic import BaseModel
import asyncio
import logging
from typing import Callable, List
from autogen_core import AgentRuntime, Component, ComponentModel
from autogen_core.models import ChatCompletionClient
from typing_extensions import Self

from backend.teams.bulter_magentic_one_orchestrator import BulteerMagenticOneOrchestrator
from backend.teams.prompts import BULTER_FINAL_ANSWER_PROMPT


#定义精英团，书房多角色工作群
class EliteGroupChatConfig(BaseModel):
    """精英团，书房多角色工作群配置类"""
    participants: List[ComponentModel]  # 参与精英团的成员组件模型列表（定义具体参与对话的角色）
    model_client: ComponentModel  # 用于对话生成的模型客户端组件模型（负责实际生成回复的大模型配置）
    termination_condition: ComponentModel | None = None  # 对话终止条件的组件模型（可选，定义对话结束的判断逻辑）
    max_turns: int | None = None  # 最大对话轮数限制（可选，不设置则无限制，防止无限循环）
    max_stalls: int  # 最大停滞次数（对话无有效进展时的终止阈值，例如连续无实质内容回复）
    final_answer_prompt: str  # 要求生成最终答案的提示文本（触发成员输出总结性结论的引导语）
    emit_team_events: bool = False  # 是否触发团队事件（默认不触发，设为True时会广播团队协作相关事件）

class EliteGroupChat(BaseGroupChat, Component[EliteGroupChatConfig]):
    """一个团队，运行由 MagenticOneOrchestrator 管理参与者的群组聊天。
    编排器负责处理对话流程，通过管理参与者之间的交互，确保高效完成任务。
    编排器基于 Magentic-One 架构，这是一个用于解决复杂任务的通用多智能体系统（请参阅下面的参考资料）。
    参数:
        participants (List[ChatAgent]): 群组聊天中的参与者。
        model_client (ChatCompletionClient): 用于生成回复的模型客户端。
        termination_condition (TerminationCondition, 可选): 群组聊天的终止条件。默认为 None。
            如果没有终止条件，群组聊天将根据编排器逻辑运行，或者直到达到最大轮数。
        max_turns (int, 可选): 群组聊天在停止前的最大轮数。默认为 20。
        max_stalls (int, 可选): 在重新规划之前允许的最大停滞次数。默认为 3。
        final_answer_prompt (str, 可选): 用于从团队对话记录中生成最终答案或回复的大语言模型提示。提供了一个默认提示（适用于 GPT-4o 类模型）。
        custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], 可选): 群组聊天中使用的自定义消息类型列表。
            如果你使用了自定义消息类型，或者你的智能体生成了自定义消息类型，则需要在此处指定。
            确保你的自定义消息类型是 :class:`~autogen_agentchat.messages.BaseAgentEvent` 或 :class:`~autogen_agentchat.messages.BaseChatMessage` 的子类。
        emit_team_events (bool, 可选): 是否通过 :meth:`BaseGroupChat.run_stream` 发出团队事件。默认为 False。
    抛出异常:
        ValueError: 如果编排逻辑中的进度分类账缺少必需的键，或者下一个发言者无效。
    """
    component_config_schema = EliteGroupChatConfig
    component_provider_override = "backend.teams.EliteGroupChat"

    def __init__(
            self,
            participants: List[ChatAgent],
            model_client: ChatCompletionClient,
            *,
            termination_condition: TerminationCondition | None = None,
            max_turns: int | None = 20,
            runtime: AgentRuntime | None = None,
            max_stalls: int = 3,
            final_answer_prompt: str = BULTER_FINAL_ANSWER_PROMPT,
            custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
            emit_team_events: bool = False,
    ):
        super().__init__(
            participants,
            group_chat_manager_name="BulteerMagenticOneOrchestrator",
            group_chat_manager_class=BulteerMagenticOneOrchestrator,
            termination_condition=termination_condition,
            max_turns=max_turns,
            runtime=runtime,
            custom_message_types=custom_message_types,
            emit_team_events=emit_team_events,
        )
        # Validate the participants.
        if len(participants) == 0:
            raise ValueError("EliteGroupChat会话必须至少有一个参与者。")
        self._model_client = model_client
        self._max_stalls = max_stalls
        self._final_answer_prompt = final_answer_prompt

    def _create_group_chat_manager_factory(
        self,
        name: str,
        group_topic_type: str,
        output_topic_type: str,
        participant_topic_types: List[str],
        participant_names: List[str],
        participant_descriptions: List[str],
        output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
        termination_condition: TerminationCondition | None,
        max_turns: int | None,
        message_factory: MessageFactory,
    ) -> Callable[[], BulteerMagenticOneOrchestrator]:
        # 定义一个工厂函数，用于创建一个新的 BulteerMagenticOneOrchestrator 实例。
        return lambda: BulteerMagenticOneOrchestrator(
            name,
            group_topic_type,
            output_topic_type,
            participant_topic_types,
            participant_names,
            participant_descriptions,
            max_turns,
            message_factory,
            self._model_client,
            self._max_stalls,
            self._final_answer_prompt,
            output_message_queue,
            termination_condition,
            self._emit_team_events,
        )

    # 从工作组实例中提取配置信息，包括参与者、模型客户端、终止条件、最大轮数、最终答案提示和是否触发团队事件。
    def _to_config(self) -> EliteGroupChatConfig:
        participants = [participant.dump_component() for participant in self._participants]
        termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
        return EliteGroupChatConfig(
            participants=participants,
            model_client=self._model_client.dump_component(),
            termination_condition=termination_condition,
            max_turns=self._max_turns,
            max_stalls=self._max_stalls,
            final_answer_prompt=self._final_answer_prompt,
            emit_team_events=self._emit_team_events,
        )

    @classmethod
    def _from_config(cls, config: EliteGroupChatConfig) -> Self:
        # 从配置信息中加载一个新的 EliteGroupChat 实例。
        logging.info(f"从EliteGroupChat配置构建一个实例对象: {config}")
        participants = [ChatAgent.load_component(participant) for participant in config.participants]
        model_client = ChatCompletionClient.load_component(config.model_client)
        termination_condition = (
            TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
        )
        return cls(
            participants,
            model_client,
            termination_condition=termination_condition,
            max_turns=config.max_turns,
            max_stalls=config.max_stalls,
            final_answer_prompt=config.final_answer_prompt,
            emit_team_events=config.emit_team_events,
        )